#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>

#define OTA_HEADER_SIZE 32
#define MAX_FIRMWARE_TYPE_LEN 10
#define VERSION_STR_LEN 32

// OTA文件头结构体
typedef struct {
    unsigned char magic1;        // 0xAA
    unsigned char magic2;        // 0x55
    unsigned char fw_count;      // 固定为1
    unsigned char fw_ver_major;  // 主版本号
    unsigned char fw_ver_minor;  // 次版本号
    unsigned char reserved[11];  // 预留字段
    unsigned int fw_size;        // 固件长度
    unsigned short fw_checksum;  // 校验和
    char fw_type[10];           // 固件类型
} ota_header_t;

// 函数声明
void show_help(void);
int pack_ota(const char *file, const char *type, const char *version);
int unpack_ota(const char *file);
int view_ota(const char *file);
int process_hex_file(const char *input, const char *output);
int process_cly_file(const char *input, const char *output);
unsigned short calc_checksum(const unsigned char *data, int len);

int main(int argc, char *argv[]) {
    int opt;
    char *input_file = NULL;
    char *fw_type = NULL;
    char *fw_version = NULL;

    while ((opt = getopt(argc, argv, "hp:u:v:")) != -1) {
        switch (opt) {
            case 'h':
                show_help();
                return 0;
            case 'p':
                if (argc < 5) {
                    printf("Error: Missing parameters for pack operation\n");
                    return -1;
                }
                input_file = optarg;
                fw_type = argv[optind];
                fw_version = argv[optind + 1];
                return pack_ota(input_file, fw_type, fw_version);
            case 'u':
                input_file = optarg;
                return unpack_ota(input_file);
            case 'v':
                input_file = optarg;
                return view_ota(input_file);
            default:
                show_help();
                return -1;
        }
    }

    show_help();
    return 0;
}

void show_help(void) {
    printf("OTA Package Tool v1.0\n");
    printf("Usage:\n");
    printf("  -h                    Show this help message\n");
    printf("  -p [file] [type] [version]  Pack firmware to OTA package\n");
    printf("  -u [file]            Unpack OTA package to binary\n");
    printf("  -v [file]            View OTA package information\n\n");
    printf("Supported file formats:\n");
    printf("  .hex - Intel Hex format\n");
    printf("  .bin - Binary format\n");
    printf("  .cly - Custom format\n\n");
    printf("Version format: a.b (a,b: 0-255)\n");
}

int pack_ota(const char *file, const char *type, const char *fw_version) {
    char *temp_bin = NULL;
    char output_file[256];
    char *dot;
    FILE *fp_in, *fp_out;
    ota_header_t header = {0};
    unsigned char *fw_data;
    long file_size;
    int major_ver, minor_ver;

    // 检查参数
    if (!file || !type || !fw_version) {
        printf("Error: Invalid parameters\n");
        return -1;
    }

    // 解析版本号
    if (sscanf(fw_version, "%d.%d", &major_ver, &minor_ver) != 2 ||
        major_ver < 0 || major_ver > 255 || minor_ver < 0 || minor_ver > 255) {
        printf("Error: Invalid version format\n");
        return -1;
    }

    // 处理输入文件
    dot = strrchr(file, '.');
    if (!dot) {
        printf("Error: Invalid input file\n");
        return -1;
    }

    // 根据文件类型进行处理
    if (strcmp(dot, ".hex") == 0) {
        temp_bin = "temp.bin";
        if (process_hex_file(file, temp_bin) != 0) {
            return -1;
        }
        file = temp_bin;
    } else if (strcmp(dot, ".cly") == 0) {
        temp_bin = "temp.bin";
        if (process_cly_file(file, temp_bin) != 0) {
            return -1;
        }
        file = temp_bin;
    }

    // 读取固件数据
    fp_in = fopen(file, "rb");
    if (!fp_in) {
        printf("Error: Cannot open input file\n");
        return -1;
    }

    fseek(fp_in, 0, SEEK_END);
    file_size = ftell(fp_in);
    fseek(fp_in, 0, SEEK_SET);

    fw_data = (unsigned char *)malloc(file_size);
    if (!fw_data) {
        fclose(fp_in);
        printf("Error: Memory allocation failed\n");
        return -1;
    }

    if (fread(fw_data, 1, file_size, fp_in) != file_size) {
        free(fw_data);
        fclose(fp_in);
        printf("Error: Failed to read firmware data\n");
        return -1;
    }
    fclose(fp_in);

    // 准备OTA头
    header.magic1 = 0xAA;
    header.magic2 = 0x55;
    header.fw_count = 1;
    header.fw_ver_major = major_ver;
    header.fw_ver_minor = minor_ver;
    header.fw_size = file_size;
    header.fw_checksum = calc_checksum(fw_data, file_size);
    strncpy(header.fw_type, type, MAX_FIRMWARE_TYPE_LEN - 1);

    // 创建输出文件
    strncpy(output_file, file, sizeof(output_file) - 1);
    dot = strrchr(output_file, '.');
    if (dot) *dot = '\0';
    strcat(output_file, "_ota.bin");

    fp_out = fopen(output_file, "wb");
    if (!fp_out) {
        free(fw_data);
        printf("Error: Cannot create output file\n");
        return -1;
    }

    // 写入数据
    fwrite(&header, 1, sizeof(header), fp_out);
    fwrite(fw_data, 1, file_size, fp_out);

    fclose(fp_out);
    free(fw_data);

    if (temp_bin) {
        remove(temp_bin);
    }

    printf("OTA package created successfully: %s\n", output_file);
    return 0;
}

int unpack_ota(const char *file) {
    FILE *fp_in, *fp_out;
    ota_header_t header;
    unsigned char *fw_data;
    char output_file[256];
    char *dot;

    fp_in = fopen(file, "rb");
    if (!fp_in) {
        printf("Error: Cannot open input file\n");
        return -1;
    }

    // 读取头信息
    if (fread(&header, 1, sizeof(header), fp_in) != sizeof(header)) {
        fclose(fp_in);
        printf("Error: Failed to read OTA header\n");
        return -1;
    }

    // 验证头信息
    if (header.magic1 != 0xAA || header.magic2 != 0x55) {
        fclose(fp_in);
        printf("Error: Invalid OTA file format\n");
        return -1;
    }

    // 分配内存并读取固件数据
    fw_data = (unsigned char *)malloc(header.fw_size);
    if (!fw_data) {
        fclose(fp_in);
        printf("Error: Memory allocation failed\n");
        return -1;
    }

    if (fread(fw_data, 1, header.fw_size, fp_in) != header.fw_size) {
        free(fw_data);
        fclose(fp_in);
        printf("Error: Failed to read firmware data\n");
        return -1;
    }

    fclose(fp_in);

    // 验证校验和
    if (header.fw_checksum != calc_checksum(fw_data, header.fw_size)) {
        free(fw_data);
        printf("Error: Checksum verification failed\n");
        return -1;
    }

    // 创建输出文件
    strncpy(output_file, file, sizeof(output_file) - 1);
    dot = strrchr(output_file, '.');
    if (dot) *dot = '\0';
    strcat(output_file, "_unpacked.bin");

    fp_out = fopen(output_file, "wb");
    if (!fp_out) {
        free(fw_data);
        printf("Error: Cannot create output file\n");
        return -1;
    }

    fwrite(fw_data, 1, header.fw_size, fp_out);
    fclose(fp_out);
    free(fw_data);

    printf("OTA package unpacked successfully: %s\n", output_file);
    return 0;
}

int view_ota(const char *file) {
    FILE *fp;
    ota_header_t header;

    fp = fopen(file, "rb");
    if (!fp) {
        printf("Error: Cannot open file\n");
        return -1;
    }

    if (fread(&header, 1, sizeof(header), fp) != sizeof(header)) {
        fclose(fp);
        printf("Error: Failed to read OTA header\n");
        return -1;
    }

    fclose(fp);

    if (header.magic1 != 0xAA || header.magic2 != 0x55) {
        printf("Error: Invalid OTA file format\n");
        return -1;
    }

    printf("OTA Package Information:\n");
    printf("Firmware Type: %s\n", header.fw_type);
    printf("Version: %d.%d\n", header.fw_ver_major, header.fw_ver_minor);
    printf("Firmware Size: %d bytes\n", header.fw_size);
    printf("Checksum: 0x%04X\n", header.fw_checksum);

    return 0;
}

unsigned short calc_checksum(const unsigned char *data, int len) {
    unsigned short sum = 0;
    int i;
    for (i = 0; i < len; i++) {
        sum += data[i];
    }
    return sum;
}

int process_hex_file(const char *input, const char *output) {
    printf("Error: HEX file processing not implemented\n");
    printf("Please use hex2bin tool to convert HEX file to BIN first\n");
    return -1;
}

int process_cly_file(const char *input, const char *output) {
    FILE *fp_in, *fp_out;
    char line[512];
    int length, i;
    unsigned char byte;
    unsigned int total_length;

    fp_in = fopen(input, "r");
    if (!fp_in) {
        printf("Error: Cannot open CLY file\n");
        return -1;
    }

    fp_out = fopen(output, "wb");
    if (!fp_out) {
        fclose(fp_in);
        printf("Error: Cannot create output file\n");
        return -1;
    }

    // 读取第一行（总长度）
    if (!fgets(line, sizeof(line), fp_in)) {
        fclose(fp_in);
        fclose(fp_out);
        printf("Error: Failed to read CLY file length\n");
        return -1;
    }

    sscanf(line, "%x", &total_length);

    // 读取数据行
    while (fgets(line, sizeof(line), fp_in)) {
        length = strlen(line);
        if (length < 2) continue;

        // 移除换行符
        if (line[length-1] == '\n') {
            line[length-1] = '\0';
            length--;
        }

        // 每两个字符转换为一个字节
        for (i = 0; i < length; i += 2) {
            sscanf(&line[i], "%2hhx", &byte);
            fwrite(&byte, 1, 1, fp_out);
        }
    }

    fclose(fp_in);
    fclose(fp_out);
    return 0;
}